package com.feidee.fd.sml.algorithm.component.ml.ps.classification

import com.feidee.fd.sml.algorithm.component.ml.MLParam
import com.tencent.angel.sona.ml.PipelineStage
import com.tencent.angel.sona.ml.classification.{AngelClassifier, CusAngelClassifier}

/**
  * @author xiongjun
  * @date 2019/12/3 15:53
  * @description
  * @reviewer
  */
case class PSLinearSVCParam(
                           override val input_pt: String,
                           override val output_pt: String,
                           override val hive_table: String,
                           override val flow_time: String,
                           override val featuresCol: String,
                           override var labelCol: String,
                           override var predictionCol: String,
                           override val modelPath: String,
                           override val metrics: Array[String],
                           // 原始结果列
                           rawPredictionCol: String,
                           // 最大迭代次数，>= 0，默认 100
                           maxIter: Int,
                           lr:Double,
                           jsonFile:String="./svm.json",
                           modelName:String="SVC"
                         ) extends MLParam {

  def this() = this(null, null, null, null, "features", "label", "prediction", null, new Array[String](0),
    "rawPrediction", 100,0.1,"./svm.json","SVC")

  override def verify(): Unit = {
    super.verify()
    require(maxIter >= 0, "param maxIter can't be negative")
  }

  override def toMap: Map[String, Any] = {
    var map = super.toMap
    map += ("rawPredictionCol" -> rawPredictionCol)
    map += ("maxIter" -> maxIter)
    map += ("lr" -> lr)
    map += ("jsonFile" -> jsonFile)
    map += ("modelName" -> modelName)
    map
  }
}

class PSLinearSVCComponent extends AbstractPSClassificationComponent[PSLinearSVCParam] {
  override def setUp(param: PSLinearSVCParam): PipelineStage = {

    val classifier = new CusAngelClassifier()
      .setNumClass(2)
      .setMaxIter(param.maxIter)
      .setLearningRate(param.lr)
      .setFeaturesCol(param.featuresCol)
      .setPredictionCol(param.predictionCol)
      .setRawPredictionCol(param.rawPredictionCol)
      .setLabelCol(param.labelCol)
      .setModelJsonFile(param.jsonFile)
    classifier
  }
}
object PSLinearSVCComponent {
  def apply(paramStr: String): Unit = {
    new PSLinearSVCComponent()(paramStr)
  }

  def main(args: Array[String]): Unit = {
    PSLinearSVCComponent(args(0))
  }

}